//
// Created by benxb on 2022/3/6.
//

#include "gtest/gtest.h"
#include <iostream>
#include <pqxx/pqxx>
#include <or_reader.h>
#include <ns_reader.h>
#include <kitti_reader.h>
#include <common/timer.h>
#include <common/ini_reader.h>

void CreateFlatTable() {
    pqxx::connection c{"postgresql://postgres:postgres@localhost/postgres"};
    pqxx::work txn{c};

    txn.exec("CREATE TABLE point_cloud( "
             "id SERIAL NOT NULL, "
             "x    REAL  NOT NULL,"
             "y    REAL  NOT NULL,"
             "z    REAL  NOT NULL,"
             "sweep INTEGER  NOT NULL)");

    txn.exec("CREATE INDEX sweep_index ON point_cloud(sweep)");
    txn.commit();
}

void DeleteFlatTable() {
    pqxx::connection c{"postgresql://postgres:postgres@localhost/postgres"};
    pqxx::work txn{c};

    txn.exec("DROP TABLE point_cloud;");
    txn.commit();
}

void WritePoints(DataReader* data_reader) {
    benxdb::Timer timer;
    timer.Start();
    CreateFlatTable();
    timer.StopAndPrint("Init time.");

    // open data file
    std::vector<float> x_arr;
    std::vector<float> y_arr;
    std::vector<float> z_arr;

    pqxx::connection c{"postgresql://postgres:postgres@localhost/postgres"};

    int scan_index = 0;
    while(data_reader->HasMoreSweep()) {
        scan_index++;
        data_reader->ReadSweep(x_arr, y_arr, z_arr);

        timer.Start();
        pqxx::work txn{c};
        auto write_stream = pqxx::stream_to::table(
            txn,
            {"point_cloud"},
            {"x", "y", "z", "sweep"}
            );

        for (int i = 0; i < x_arr.size(); i++) {
            write_stream << std::make_tuple(x_arr[i], y_arr[i], z_arr[i], scan_index);
        }
        write_stream.complete();
        timer.Pause();

        x_arr.clear();
        y_arr.clear();
        z_arr.clear();
    }

    timer.StopAndPrint("Load time.");
}

TEST(PgFlat, LoadNSData) {
    INIReader reader("config.ini");
    std::string data_dir = reader.Get("data", "data_dir", "");
    std::string ns01 = reader.Get("nuScenes", "01", "");
    DataReader* data_reader = new NSReader(data_dir + ns01);
    DeleteFlatTable();
    WritePoints(data_reader);
    delete data_reader;
}

TEST(PgFlat, LoadOR01Data) {
    INIReader reader("config.ini");
    std::string data_dir = reader.Get("data", "data_dir", "");
    std::string or01 = reader.Get("oxford", "01", "");
    DataReader* data_reader = new ORReader(data_dir + or01);
    DeleteFlatTable();
    WritePoints(data_reader);
    delete data_reader;
}

TEST(PgFlat, LoadOR02Data) {
    INIReader reader("config.ini");
    std::string data_dir = reader.Get("data", "data_dir", "");
    std::string or02 = reader.Get("oxford", "02", "");
    DataReader* data_reader = new ORReader(data_dir + or02);
    DeleteFlatTable();
    WritePoints(data_reader);
    delete data_reader;
}

TEST(PgFlat, LoadKT01Data) {
    INIReader reader("config.ini");
    std::string data_dir = reader.Get("data", "data_dir", "");
    std::string kt01 = reader.Get("kitti", "01", "");
    DataReader* data_reader = new ORReader(data_dir + kt01);
    DeleteFlatTable();
    WritePoints(data_reader);
    delete data_reader;
}

TEST(PgFlat, LoadKT02Data) {
    INIReader reader("config.ini");
    std::string data_dir = reader.Get("data", "data_dir", "");
    std::string kt02 = reader.Get("kitti", "02", "");
    DataReader* data_reader = new ORReader(data_dir + kt02);
    DeleteFlatTable();
    WritePoints(data_reader);
    delete data_reader;
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc,argv);
    return RUN_ALL_TESTS();
}